#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
交易策略模块
包含各种量化交易策略的实现
"""

import pandas as pd


def ma_strategy(data, short_window=5, long_window=20):
    """
    双均线交易策略
    当短期均线向上穿越长期均线时买入，当短期均线向下穿越长期均线时卖出
    
    参数:
    data: 股票数据DataFrame
    short_window: 短期均线窗口
    long_window: 长期均线窗口
    
    返回:
    DataFrame: 包含信号的数据
    """
    # 复制数据以避免修改原始数据
    signals = data.copy()
    
    # 计算短期和长期移动平均线
    signals['short_ma'] = data['close'].rolling(window=short_window, min_periods=1).mean()
    signals['long_ma'] = data['close'].rolling(window=long_window, min_periods=1).mean()
    
    # 初始化信号列
    signals['signal'] = 0.0
    
    # 当短期均线上穿长期均线时，产生买入信号(值为1.0)
    # 当短期均线下穿长期均线时，产生卖出信号(值为-1.0)
    # 使用loc方法避免pandas警告
    signals.loc[short_window:, 'signal'] = (
        signals['short_ma'][short_window:] > signals['long_ma'][short_window:]
    ).astype(float)
    
    # 计算信号的变化，1表示买入，-1表示卖出
    signals['positions'] = signals['signal'].diff()
    
    return signals


def evaluate_strategy(signals):
    """
    评估策略表现
    
    参数:
    signals: 包含交易信号的数据
    
    返回:
    dict: 策略评估结果
    """
    # 计算总的交易信号数量
    total_signals = len(signals[signals['positions'] != 0])
    buy_signals = len(signals[signals['positions'] == 1])
    sell_signals = len(signals[signals['positions'] == -1])
    
    return {
        'total_signals': total_signals,
        'buy_signals': buy_signals,
        'sell_signals': sell_signals
    }